"""
时域特征提取

实现了时域特征提取功能，包括：
- 振幅特征
- 能量特征
- 过零率
- RMS特征
- 峰峰值
"""

import numpy as np
from typing import Dict, List
from loguru import logger


class TimeDomainFeatures:
    """时域特征提取器"""
    
    def __init__(self, features: List[str] = None):
        """
        初始化时域特征提取器
        
        Args:
            features: 要提取的特征列表
        """
        if features is None:
            features = ["amplitude", "energy", "zero_crossing_rate", "rms", "peak_to_peak"]
        
        self.features = features
        logger.info(f"初始化时域特征提取器，特征: {features}")
    
    def extract_features(self, signal_data: np.ndarray) -> Dict[str, float]:
        """
        提取时域特征
        
        Args:
            signal_data: 输入信号
            
        Returns:
            特征字典
        """
        feature_dict = {}
        
        for feature in self.features:
            if feature == "amplitude":
                feature_dict[feature] = self._extract_amplitude(signal_data)
            elif feature == "energy":
                feature_dict[feature] = self._extract_energy(signal_data)
            elif feature == "zero_crossing_rate":
                feature_dict[feature] = self._extract_zero_crossing_rate(signal_data)
            elif feature == "rms":
                feature_dict[feature] = self._extract_rms(signal_data)
            elif feature == "peak_to_peak":
                feature_dict[feature] = self._extract_peak_to_peak(signal_data)
            else:
                logger.warning(f"未知的时域特征: {feature}")
        
        return feature_dict
    
    def _extract_amplitude(self, signal_data: np.ndarray) -> float:
        """
        提取振幅特征
        
        Args:
            signal_data: 输入信号
            
        Returns:
            振幅特征值
        """
        return np.max(np.abs(signal_data))
    
    def _extract_energy(self, signal_data: np.ndarray) -> float:
        """
        提取能量特征
        
        Args:
            signal_data: 输入信号
            
        Returns:
            能量特征值
        """
        return np.sum(signal_data ** 2)
    
    def _extract_zero_crossing_rate(self, signal_data: np.ndarray) -> float:
        """
        提取过零率特征
        
        Args:
            signal_data: 输入信号
            
        Returns:
            过零率特征值
        """
        signs = np.sign(signal_data)
        zero_crossings = np.sum(np.abs(signs[1:] - signs[:-1])) / 2
        return zero_crossings / len(signal_data)
    
    def _extract_rms(self, signal_data: np.ndarray) -> float:
        """
        提取RMS特征
        
        Args:
            signal_data: 输入信号
            
        Returns:
            RMS特征值
        """
        return np.sqrt(np.mean(signal_data ** 2))
    
    def _extract_peak_to_peak(self, signal_data: np.ndarray) -> float:
        """
        提取峰峰值特征
        
        Args:
            signal_data: 输入信号
            
        Returns:
            峰峰值特征值
        """
        return np.max(signal_data) - np.min(signal_data)
    
    def extract_all_features(self, signal_data: np.ndarray) -> Dict[str, float]:
        """
        提取所有时域特征
        
        Args:
            signal_data: 输入信号
            
        Returns:
            所有特征字典
        """
        all_features = [
            "amplitude", "energy", "zero_crossing_rate", "rms", "peak_to_peak"
        ]
        
        feature_dict = {}
        for feature in all_features:
            if feature == "amplitude":
                feature_dict[feature] = self._extract_amplitude(signal_data)
            elif feature == "energy":
                feature_dict[feature] = self._extract_energy(signal_data)
            elif feature == "zero_crossing_rate":
                feature_dict[feature] = self._extract_zero_crossing_rate(signal_data)
            elif feature == "rms":
                feature_dict[feature] = self._extract_rms(signal_data)
            elif feature == "peak_to_peak":
                feature_dict[feature] = self._extract_peak_to_peak(signal_data)
        
        return feature_dict
    
    def get_feature_names(self) -> List[str]:
        """
        获取特征名称列表
        
        Returns:
            特征名称列表
        """
        return self.features.copy()
    
    def __str__(self) -> str:
        return f"TimeDomainFeatures(features={self.features})"
    
    def __repr__(self) -> str:
        return self.__str__() 